from ultralytics import YOLO

model = YOLO("../../yolov8n.pt")  # 下载或加载预训练模型（建议用于训练）

if __name__ == '__main__':
    # 使用模型
    # data指定数据集配置文件的路径
    # epochs指定训练的轮数
    # workers指定用于数据加载的子进程数，0表示使用主进程
    # mixup指定是否使用mixup数据增强，0.5表示50%的概率使用
    # mosaic指定是否使用mosaic数据增强，0.5表示50%的概率使用
    # device指定训练使用的设备，0表示使用第一个GPU
    model.train(data="dataset/am.yaml", epochs=10, workers=0, mixup=0.5, mosaic=0.5)  # 训练模型
    # metrics = model.val()  # 在验证集上评估模型性能
    # 将模型导出为 ONNX 格式，
    # 作用：ONNX 是一种开放格式，用于表示深度学习模型。它旨在提供一种标准格式，使得不同框架之间可以互相转换模型。
    # opset指定ONNX的opset版本
    # 返回导出是否成功
    success = model.export(format="onnx", opset=12)
    # 对指定图像进行预测
    # 返回一个结果列表，每个元素是一个Prediction对象
    results = model("0049.jpg")  # 对图像进行预测
    # Process results list
    for result in results:
        # 获取边界框输出
        boxes = result.boxes
        # 获取分割掩码输出
        masks = result.masks
        # 获取姿态关键点输出
        keypoints = result.keypoints
        # 获取分类概率输出
        probs = result.probs
        # 显示预测结果到屏幕
        result.show()
